// https://leetcode.cn/problems/binary-tree-pruning/

// 题干：给你二叉树的根结点 root ，此外树的每个结点的值要么是 0 ，要么是 1 。
//       返回移除了所有不包含 1 的子树的原二叉树。
//       节点 node 的子树为 node 本身加上所有 node 的后代。

// 示例：输入：root = [1,null,0,0,1]
//       输出：[1,null,0,null,1](层序遍历)
//             [1,0,0,1](前序遍历)

// 碎语：因为是要删子树，考虑后序遍历，返回的是节点指针，函数类型确定，仅仅需要传入一个根节点
//       传入参数确定，返回情况分类讨论，如果节点的值为0，也一并删除，即返回nullptr,若非0，则返回该节点

#include <bits/stdc++.h>
using namespace std;

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

class Solution
{
public:
    TreeNode* pruneTree(TreeNode* root)
    {
        if (root == nullptr) return nullptr;

        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);

        if (root->left == nullptr && root->right == nullptr && root->val == 0){
            delete root;
            root = nullptr;
        }

        return root;
    }

    void preOrder(TreeNode* root)
    {
        if (root == nullptr) return ;

        cout << root->val << " ";
        preOrder(root->left);
        preOrder(root->right);
    }
};

int main()
{
    //root = [1,null,0,0,1]
    Solution sol;

    auto root = new TreeNode(1);
    root->right = new TreeNode(0);
    root->right->left = new TreeNode(0);
    root->right->right = new TreeNode(1);

    // 前序遍历来验证
    sol.preOrder(root);

    return 0;
}

